/*
 * Copyright(C) 2020. Huawei Technologies Co.,Ltd. All rights reserved.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#ifndef TOPK_OP_INCLUDED
#define TOPK_OP_INCLUDED

#include <arm_fp16.h>
#include <ascenddaemon/utils/AscendTensor.h>

namespace ascend {
class TopkOp {
public:
    TopkOp();
    virtual ~TopkOp();

    bool exec(AscendTensor<float16_t, DIMS_2>& distance,
              AscendTensor<uint32_t, DIMS_2>& indices,
              AscendTensor<float16_t, DIMS_2>& topkDistance,
              AscendTensor<uint32_t, DIMS_2>& topkIndices,
              uint32_t indexOffset = 0);

    bool exec(AscendTensor<float16_t, DIMS_1> &distance,
                  AscendTensor<uint32_t, DIMS_1> &indices,
                  AscendTensor<float16_t, DIMS_1> &topkDistance,
                  AscendTensor<uint32_t, DIMS_1> &topkIndices,
                  uint32_t realSize,
                  uint32_t indexOffset);

    bool exec(AscendTensor<float16_t, DIMS_1>& distance,
              AscendTensor<uint32_t, DIMS_1>& indices,
              AscendTensor<float16_t, DIMS_1>& topkDistance,
              AscendTensor<uint32_t, DIMS_1>& topkIndices);

    void reorder(AscendTensor<float16_t, DIMS_2>& topkDistance,
                 AscendTensor<uint32_t, DIMS_2>& topkIndices);

    void reorder(AscendTensor<float16_t, DIMS_1>& topkDistance,
                 AscendTensor<uint32_t, DIMS_1>& topkIndices);

private:
    bool checkParams(AscendTensor<float16_t, DIMS_2>& distance,
                     AscendTensor<uint32_t, DIMS_2>& indices,
                     AscendTensor<float16_t, DIMS_2>& topkDistance,
                     AscendTensor<uint32_t, DIMS_2>& topkIndices);

    bool checkParams(AscendTensor<float16_t, DIMS_1>& distance,
                     AscendTensor<uint32_t, DIMS_1>& indices,
                     AscendTensor<float16_t, DIMS_1>& topkDistance,
                     AscendTensor<uint32_t, DIMS_1>& topkIndices);

    void minValuePush(size_t size, float16_t* heapDist,
                      uint32_t* heapId, float16_t pushDist, uint32_t pushId);

    void maxValuePop(size_t k, float16_t* heapDist, uint32_t* heapId);
};
}  // namespace ascend

#endif
